package com.changan;

import ai.onnxruntime.OnnxTensor;
import ai.onnxruntime.OrtEnvironment;
import ai.onnxruntime.OrtException;
import ai.onnxruntime.OrtSession;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.MutationResult;
import io.milvus.grpc.SearchResults;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.response.SearchResultsWrapper;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ArrayUtils;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.IOException;
import java.nio.FloatBuffer;
import java.util.*;

@Slf4j
@Service
public class DefaultMilvusService implements MilvusService, DisposableBean {

    @Autowired
    private MilvusServiceClient milvusClient;

    @Value("${onnx.model.path:resnet50-v1-12.onnx}")
    private String modelPath;

    private OrtEnvironment env;
    private OrtSession session;

    @PostConstruct
    public void initialize() {
        try {
            env = OrtEnvironment.getEnvironment();
            // 使用 Spring 的资源加载，避免 jar 包内找不到模型文件
            Resource resource = new ClassPathResource(modelPath);
            session = env.createSession(
                    resource.getFile().getAbsolutePath(),
                    new OrtSession.SessionOptions()
            );
            log.info("ONNX 环境初始化完成，模型路径: {}", modelPath);
        } catch (Exception e) {
            log.error("ONNX 初始化失败: {}", e.getMessage(), e);
            throw new IllegalStateException("初始化 ONNX 模型失败", e);
        }
    }

    /**
     * 提取图片的特征向量
     */
    public float[] extractFeature(String imagePath) {
        log.info("开始提取图片特征: {}", imagePath);
        try {
            float[] input = ImageUtils.preprocess(new File(imagePath));
            try (OnnxTensor tensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(input),
                    new long[]{1, 3, 224, 224});
                 OrtSession.Result output = session.run(Collections.singletonMap("data", tensor))) {

                float[][] result = (float[][]) output.get(0).getValue();
                log.info("图片 {} 特征提取完成，特征向量维度: {}", imagePath, result[0].length);
                return result[0];
            }
        } catch (OrtException e) {
            log.error("ONNX 推理失败: {}", e.getMessage(), e);
            throw new RuntimeException("ONNX 推理失败: " + imagePath, e);
        } catch (IOException e) {
            log.error("图像文件读取失败: {}", imagePath, e);
            throw new RuntimeException("图像文件读取失败: " + imagePath, e);
        } catch (Exception e) {
            log.error("提取图片特征时出错: {}", imagePath, e);
            throw new RuntimeException("未知错误: " + imagePath, e);
        }
    }

    /**
     * 将特征向量插入到 Milvus
     */
    public void insertToMilvus(float[] vector, String imageUrl) {
        String uuid = UUID.randomUUID().toString();
        List<InsertParam.Field> fields = new ArrayList<>();
        fields.add(new InsertParam.Field("id", Collections.singletonList(uuid)));
        fields.add(new InsertParam.Field("vector", Collections.singletonList(Arrays.asList(ArrayUtils.toObject(vector)))));
        fields.add(new InsertParam.Field("url", Collections.singletonList(imageUrl)));

        InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName("image_collection")
                .withFields(fields)
                .build();

        log.info("开始插入数据到 Milvus，ID: {}, 图片URL: {}", uuid, imageUrl);
        R<MutationResult> response = milvusClient.insert(insertParam);
        if (response.getStatus() != R.Status.Success.getCode()) {
            log.error("数据插入失败: {}", response.getMessage());
            throw new RuntimeException("Milvus 插入失败: " + response.getMessage());
        }
        log.info("数据插入完成，ID: {}", uuid);
    }

    /**
     * 在 Milvus 中搜索相似图片
     */
    @Override
    public List<String> searchInMilvus(float[] vector) {
        log.info("开始在 Milvus 中搜索相似图片，向量维度: {}", vector.length);

        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName("image_collection")
                .withMetricType(MetricType.COSINE)
                .withOutFields(Collections.singletonList("url"))
                .withTopK(10)
                .withVectors(Collections.singletonList(Arrays.asList(ArrayUtils.toObject(vector))))
                .withVectorFieldName("vector")
                .build();

        R<SearchResults> response = milvusClient.search(searchParam);
        if (response.getStatus() != R.Status.Success.getCode()) {
            log.error("搜索失败: {}", response.getMessage());
            throw new RuntimeException("Milvus 搜索失败: " + response.getMessage());
        }

        SearchResultsWrapper wrapper = new SearchResultsWrapper(response.getData().getResults());
        List<SearchResultsWrapper.IDScore> scores = wrapper.getIDScore(0);

        List<String> imageUrls = new ArrayList<>();
        for (SearchResultsWrapper.IDScore score : scores) {
            Object urlObj = score.get("url");
            if (urlObj != null) {
                imageUrls.add(urlObj.toString());
            }
        }

        log.info("搜索完成，找到 {} 个相似图片", imageUrls.size());
        return imageUrls;
    }

    /**
     * Bean 销毁时清理资源
     */
    @Override
    public void destroy() {
        cleanup();
    }

    public void cleanup() {
        try {
            if (session != null) {
                session.close();
                log.info("ONNX Session 已关闭");
            }
            // OrtEnvironment 是全局单例，一般不需要主动关闭
        } catch (OrtException e) {
            log.error("关闭 ONNX 资源时出错: {}", e.getMessage(), e);
        }
    }
}
